#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <tss2/tss2_rc.h>
#include <tss2/tss2_sys.h>
#include <tss2/tss2_tctildr.h>
#include <tss2/tss2_tpm2_types.h>
#include <unistd.h>

#ifndef TCTI_NAME
#define TCTI_NAME "mssim"
#endif

TSS2_TCTI_CONTEXT* tcti_context = NULL;
TSS2_SYS_CONTEXT* sys_context = NULL;

// Record handles for cleanup
typedef struct HandleRecorder {
    TPM2_HANDLE handles[128];
    uint16_t count;
} HandleRecorder;

static HandleRecorder handle_recorder = {
    .count = 0, .handles = {TPM2_RH_NULL} // Initialize with NULL handle
};

// Check if the return code indicates success, clean up if not
#define CHECK_PASSED(rc)                                                                           \
    {                                                                                              \
        if ((rc) != TPM2_RC_SUCCESS) {                                                             \
            printf("FAILED! [%s:%u] %s\n", __FUNCTION__, __LINE__, Tss2_RC_Decode(rc));            \
            cleanup(rc);                                                                           \
        } else {                                                                                   \
            printf("PASSED! [%s:%u]\n", __FUNCTION__, __LINE__);                                   \
        }                                                                                          \
    }

#define CHUNK_SIZE 1024

// Initialize TCTI Context
TSS2_RC init_tcti_context(TSS2_TCTI_CONTEXT** tcti_context) {
    TSS2_RC rc = Tss2_TctiLdr_Initialize(TCTI_NAME, tcti_context);
    return rc;
}

// Initialize SAPI Context
TSS2_RC init_sys_context(TSS2_TCTI_CONTEXT* tcti_context, TSS2_SYS_CONTEXT** sys_context) {
    if (tcti_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    size_t sys_context_size = Tss2_Sys_GetContextSize(0);
    *sys_context = (TSS2_SYS_CONTEXT*)malloc(sys_context_size);
    if (*sys_context == NULL) {
        return TSS2_BASE_RC_MEMORY;
    }

    TSS2_ABI_VERSION abi_version = TSS2_ABI_VERSION_CURRENT;
    TSS2_RC rc = Tss2_Sys_Initialize(*sys_context, sys_context_size, tcti_context, &abi_version);
    if (rc != TSS2_RC_SUCCESS) {
        free(*sys_context);
        *sys_context = NULL;
    }

    return rc;
}

// Startup TPM
TSS2_RC startup_tpm(TSS2_SYS_CONTEXT* sys_context) {
    if (sys_context == NULL) {
        return TSS2_SYS_RC_BAD_REFERENCE;
    }

    TSS2_RC rc = Tss2_Sys_GetCapability(sys_context, NULL, TPM2_CAP_TPM_PROPERTIES,
                                        TPM2_PT_MANUFACTURER, 1, NULL, NULL, NULL);
    if (rc == TPM2_RC_INITIALIZE) {
        rc = Tss2_Sys_Startup(sys_context, TPM2_SU_CLEAR);
    }

    return rc;
}

// Cleanup resources
static void cleanup(int rc) {
    for (uint16_t i = 0; i < handle_recorder.count; i++) {
        if (handle_recorder.handles[i] != TPM2_RH_NULL) {
            TSS2_RC rc_flush = Tss2_Sys_FlushContext(sys_context, handle_recorder.handles[i]);
            if (rc_flush != TSS2_RC_SUCCESS) {
                printf("FlushContext Failed for handle 0x%08x\n", handle_recorder.handles[i]);
            }
        }
    }
    if (tcti_context != NULL) {
        Tss2_TctiLdr_Finalize(&tcti_context);
        free(tcti_context);
    }
    if (sys_context != NULL) {
        Tss2_Sys_Finalize(sys_context);
        free(sys_context);
    }
    exit(rc != TSS2_RC_SUCCESS);
    return;
}

// Create a key for HMAC
TSS2_RC create_key(TSS2_SYS_CONTEXT* sys_context, TPM2_HANDLE* key_handle) {
    // CreatePrimary Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_array = {.count = 1, .auths = {{.sessionHandle = TPM2_RH_PW}}};
    TPM2B_PUBLIC in_public = {
        .size = 0,
        .publicArea = {
            .type = TPM2_ALG_KEYEDHASH,
            .nameAlg = TPM2_ALG_SHA256,
            .objectAttributes = TPMA_OBJECT_USERWITHAUTH | TPMA_OBJECT_SIGN_ENCRYPT |
                                TPMA_OBJECT_SENSITIVEDATAORIGIN,
            .parameters.keyedHashDetail.scheme = {.scheme = TPM2_ALG_HMAC,
                                                  .details.hmac.hashAlg = TPM2_ALG_SHA256},
        }};
    TPM2B_SENSITIVE_CREATE in_sensitive = {0};
    TPM2B_DATA outside_info = {
        .size = 0,
        .buffer = {},
    };
    TPML_PCR_SELECTION creation_pcr = {
        .count = 0,
    };

    // CreatePrimary Output Parameters
    TPM2B_PUBLIC out_public = {0};
    TPM2B_CREATION_DATA creation_data = {0};
    TPM2B_DIGEST creation_hash = {.size = sizeof(creation_hash.buffer)};
    TPMT_TK_CREATION creation_ticket = {0};
    TPM2B_NAME name = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_array = {0};

    TSS2_RC rc = Tss2_Sys_CreatePrimary(sys_context, TPM2_RH_OWNER, &cmd_auths_array, &in_sensitive,
                                        &in_public, &outside_info, &creation_pcr, key_handle,
                                        &out_public, &creation_data, &creation_hash,
                                        &creation_ticket, &name, &rsp_auths_array);

    return rc;
}

// Sequence HMAC
TSS2_RC sequence_hmac(TSS2_SYS_CONTEXT* sys_context, TPM2_HANDLE key_handle, FILE* fp,
                      TPMI_DH_OBJECT* sequence_handle, TPM2B_DIGEST* out_hmac) {
    // HMAC_Start Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_start = {.count = 1, .auths = {{.sessionHandle = TPM2_RH_PW}}};
    TPM2B_AUTH auth = {0};
    TPMI_ALG_HASH hash_alg = TPM2_ALG_SHA256;

    // HMAC_Start Output Parameters
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_start = {0};

    TSS2_RC rc = Tss2_Sys_HMAC_Start(sys_context, key_handle, &cmd_auths_start, &auth, hash_alg,
                                     sequence_handle, &rsp_auths_start);
    if (rc != TSS2_RC_SUCCESS) {
        return rc;
    }
    handle_recorder.handles[handle_recorder.count++] =
        *sequence_handle; // Record the sequence handle

    // SequenceUpdate Input Parameters
    TSS2L_SYS_AUTH_COMMAND cmd_auths_update = {.count = 1,
                                               .auths = {{.sessionHandle = TPM2_RH_PW}}};

    // SequenceUpdate Output Parameters
    TPM2B_MAX_BUFFER chunk = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_update = {0};

    // Read file in chunks and update the sequence
    size_t bytes_read;
    while ((bytes_read = fread(chunk.buffer, 1, CHUNK_SIZE, fp)) > 0) {
        chunk.size = (UINT16)bytes_read;
        rc = Tss2_Sys_SequenceUpdate(sys_context, *sequence_handle, &cmd_auths_update, &chunk,
                                     &rsp_auths_update);
        if (rc != TSS2_RC_SUCCESS) {
            return rc;
        }
    }

    // SequenceComplete Input Parameters
    chunk.size = 0;
    TSS2L_SYS_AUTH_COMMAND cmd_auths_complete = {.count = 1,
                                                 .auths = {{.sessionHandle = TPM2_RH_PW}}};
    TPMI_RH_HIERARCHY hierarchy = TPM2_RH_NULL;

    // SequenceComplete Output Parameters
    TPMT_TK_HASHCHECK validation = {0};
    TSS2L_SYS_AUTH_RESPONSE rsp_auths_complete = {0};

    rc = Tss2_Sys_SequenceComplete(sys_context, *sequence_handle, &cmd_auths_complete, &chunk,
                                   hierarchy, out_hmac, &validation, &rsp_auths_complete);
    handle_recorder.count--; // Remove the sequence handle from the recorder
    return rc;
}

int main(int argc, char* argv[]) {

    TSS2_RC rc;

    // Initialize TCTI context
    rc = init_tcti_context(&tcti_context);
    CHECK_PASSED(rc);

    // Initialize SAPI context
    rc = init_sys_context(tcti_context, &sys_context);
    CHECK_PASSED(rc);

    // Startup TPM
    rc = startup_tpm(sys_context);
    CHECK_PASSED(rc);

    // Create a key
    TPM2_HANDLE key_handle;
    rc = create_key(sys_context, &key_handle);
    CHECK_PASSED(rc);
    handle_recorder.handles[handle_recorder.count++] = key_handle; // Record the key handle

    // HMAC
    FILE* fp = fopen("test.bin", "rb");
    if (fp == NULL) {
        printf("Failed to open file for HMAC\n");
        cleanup(TSS2_BASE_RC_GENERAL_FAILURE);
    }
    TPMI_DH_OBJECT sequence_handle;
    TPM2B_DIGEST out_hmac = {0};
    rc = sequence_hmac(sys_context, key_handle, fp, &sequence_handle, &out_hmac);
    CHECK_PASSED(rc);
    printf("HMAC: ");
    for (size_t i = 0; i < out_hmac.size; i++)
        printf("%02x", out_hmac.buffer[i]);
    printf("\n");

    // Cleanup resources
    cleanup(rc);

    return 0;
}
